import datetime
import warnings

import joblib
import matplotlib.pyplot as plt
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

from src.zzm.feature.feature import feature_proprecessing
from src.zzm.train.train import BrainPredictModel
from src.zzm.utils.common import data_preprocessing, plt_fig
from src.zzm.utils.log import Logger
from src.zzm.config import config

warnings.filterwarnings('ignore', category=UserWarning)

plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['font.size'] = 15


def model_predict(data, logger):
    # 1、特征预处理
    train_data = feature_proprecessing(data)

    # 2、获取 标签列
    y = train_data.pop('Attrition')
    x = train_data[config.FEATURE_NAMES]

    # 加载 标准化 模型
    scaler = joblib.load('../model/scaler_zzm.pkl')
    x = scaler.transform(x)

    # 加载随机森林模型
    model1 = joblib.load('../model/model1_zzm.pkl')
    y_pred = model1.predict_proba(x)[:, 1]

    # 将源特征数据和模型1预测之后的数据合并，再去训练
    # x_combined = np.hstack([x, x_pred_proba])
    # model2 = joblib.load('../model/model2_zzm.pkl')
    # y_pred = model2.predict_proba(x_combined)[:, 1]

    # 5、绘制ROC曲线
    plt_fig(y, y_pred)
    logger.info(f'测试集AUC:{roc_auc_score(y, y_pred)}')
    print(f'测试集AUC:{roc_auc_score(y, y_pred)}')


if __name__ == '__main__':
    pm = BrainPredictModel('../../../data/raw/test2.csv', 'test_')
    # feature_engineering(pm.data, pm.log)
    model_predict(pm.data, pm.log)
